package Lee_code;

import java.util.*;

/**
 * @author blh
 * @date 2024年10月31日 9:53
 * 查找和最小的 K 对数字
 */
public class Lee_373_1031 {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 使用优先队列模拟堆
        PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(o -> o[0]));
        List<List<Integer>> ans = new ArrayList<>();
        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            // 至多有k个
            // nums1[i] + nums2[0] 是元素，i是nums1的下标， 0是nums2的下标
            pq.add(new int[]{nums1[i] + nums2[0], i, 0});
        }
        while (ans.size() < k) {
            int[] p = pq.poll();
            int i = p[1];
            int j = p[2];
            ans.add(Arrays.asList(nums1[i], nums2[j]));
            if (j + 1 < nums2.length) {
                pq.add(new int[]{nums1[i] + nums2[j + 1], i, j + 1});
            }
        }
        return ans;
    }
}
